#!/usr/bin/env python3
# -*- coding: utf-8 -*-

''' 
    date: 2021-06-08
    author: baole 

    【功能1：从网易获取单个股票的日k数据】
    1. 通过命令行参数--single获取股票代码
    2. 从网易行情页面抓取该股票的日k数据
    网址如下：
    http://quotes.money.163.com/service/chddata.html?code=%s&start=%s&end=%s&fields=TOPEN;HIGH;LOW;TCLOSE;VOTURNOVER;VATURNOVER

        字段：
        TOPEN： 今开，开盘价
        HIGH： 最高价
        LOW： 最低价
        TCLOSE： 今收，收盘价
        LCLOSE： 昨收，昨收盘价
        TURNOVER： 换手
        VOTURNOVER：成交量
        VATURNOVER：成交金额
        CHG：涨跌额
        PCHG：涨跌幅
        TCAP：总市值
        MCAP：流通市值
    3. 保存到系统数据目录
    
    

    【功能2：从网易获取所有股票的日k数据】
    1. 通过命令行参数 --all和--single来区分获取一个股票还是所有股票
    2. 玉获取单个股票后做的操作相似

    【功能3：每日更新日k数据】
    1. 通过参数--daily 来获取当日的数据
    2. 当日的数据做处理后添加到原先的数据中最后一行。

    4. 处理数据，反转排序，使得时间从远到近，第一行是最早的时间，最后一行是最晚的时间。去掉列标题，
        数据仅仅包括时间、ohlc、量能、金额这7列即可。
    
    
'''

import os
import sys
import requests
import json
import re
import time
import random


# from quantway.lib.utils import read_jsonparam, get_module_name, get_codelist, split_list_by_num
from quantway.common.settings import ROOT_DATA_CSV_DAY
from quantway.lib.utils import get_codelist

def query_day_data(stock_code, startdate, enddate):
    ''' 处理url
        网址如下：
        http://quotes.money.163.com/service/chddata.html?code=%s&start=%s&end=%s&fields=TOPEN;HIGH;LOW;TCLOSE;VOTURNOVER;VATURNOVER

        字段：
        TOPEN： 今开，开盘价
        HIGH： 最高价
        LOW： 最低价
        TCLOSE： 今收，收盘价
        LCLOSE： 昨收，昨收盘价
        TURNOVER： 还手
        VOTURNOVER：成交量
        VATURNOVER：成交金额
        CHG：涨跌额
        PCHG：涨跌幅
        TCAP：总市值
        MCAP：流通市值
    '''
    chddata_code = stock_code.replace("sh","0") if stock_code.startswith("sh") else stock_code.replace("sz","1")
    url = "code=%s" % chddata_code
    if startdate != "":
        url += "&start=%s" % startdate
    if enddate != "":
        url += "&end=%s" % enddate
    # 数据只需要7列，日期、开盘价、最高价、最低价、收盘价、成交量、金额
    url += "&fields=TOPEN;HIGH;LOW;TCLOSE;VOTURNOVER;VATURNOVER"   
    r = query_day_data_from_163(url)
    return r.text


def query_day_data_from_163(url):
    # 对于失败的request,我们会尝试5次。
    attempt = 0
    while attempt<5:
        try:
            r = requests.get('http://quotes.money.163.com/service/chddata.html?%s' % (url))
            return r                
        except Exception as e:
            print(e)
            time.sleep(60)
        attempt+=1
        print("获取%s的日k数据超时，重试%d " % (stock_code, attempt) )

    return None


def process_data(data):
    ''' 进行处理
     去掉列标题，去掉最后一行
     顺序反转，最近的日期在最后
     去掉代码和名字,
     去除停牌的行
    '''
    li_temp = data.strip().split(os.linesep)
    # 去除列标题
    li_temp = li_temp[1:]
    li_temp.reverse()
    li_lines = list()
    for line in li_temp:
        # 排除停牌的行
        if "0.0,0.0,0.0" not in line:
            # 删除证券代码和名称
            li_lines.append(re.sub(r"'\d{6},[^,]+,", "", line))
    return os.linesep.join(li_lines)

    


    
def to_csv(stock_code, data, append=False):
    '''
        根据是否是append，选择全新创建文件保存还是添加数据到文件末尾。
    '''
    if not append:
        with open(os.path.join(ROOT_DATA_CSV_DAY, stock_code + ".csv"),"w", newline="") as f:
            f.write(data)   
        print("保存文件%s" % os.path.join(ROOT_DATA_CSV_DAY, stock_code + ".csv"))
    else:
        with open(os.path.join(ROOT_DATA_CSV_DAY, stock_code + ".csv"),"a", newline="") as f:
            f.write(data)   
        print("添加数据到文件%s" % os.path.join(ROOT_DATA_CSV_DAY, stock_code + ".csv"))


    


if __name__ == '__main__':
    '''
        脚本功能： 获取证券日k线数据，保存为csv文件
        用法：
            python query_data_day_163.py  --single sz300171 --startdate 20200101 --enddate 20210608

        参数：
            --single  证券代码   表示处理单一证券
            --all               表示处理所有证券
            这两个参数互斥。证券代码以sh，sz开头来区分沪市，深市证券
            --startdate  日期    起始日期
            --enddate    日期    结束日期
            日期格式 YYYYMMDD

    '''
    if "--startdate" in sys.argv:
        startdate = sys.argv[sys.argv.index("--startdate") + 1]
    else:
        startdate = ""
    if "--enddate" in sys.argv:
        enddate = sys.argv[sys.argv.index("--enddate") + 1]
    else:
        enddate = ""


    if "--single" in sys.argv:
        stock_code = sys.argv[sys.argv.index("--single") + 1]
        # 爬取数据
        data = query_day_data(stock_code, startdate, enddate)
        
        # 数据处理
        data = process_data(data)

        # 保存到csv文件中去
        if "--daily" not in sys.argv:
            to_csv(stock_code, data)
        else:
            to_csv(stock_code, data, True)

    elif "--all" in sys.argv:

        li_stock_code = get_codelist("all_stockcode")  
        for stock_code in li_stock_code:
            
            print("获取{}的日k数据".format(stock_code))
            data = query_day_data(stock_code, startdate, enddate) 
            if data is None:
                print("获取{}的日k数据失败".format(stock_code))
            if "--daily" not in sys.argv:
                to_csv(stock_code, data)
            else:
                to_csv(stock_code, data, True)
        

    
